import argparse
import gym
import numpy as np
import os
import torch
from torch.utils.tensorboard import SummaryWriter
import BCQ
import DDPG
import utils


# Handles interactions with the environment, i.e. train behavioral or generate buffer
def interact_with_environment(env, state_dim, action_dim, max_action, device, args):
    print("Entry process: interact_with_environment")
    # For saving files
    setting = f"{args.env}_{args.seed}"
    buffer_name = f"{args.buffer_name}_{setting}"
    summary_writer = SummaryWriter(log_dir='./runs')

    # Initialize and load policy
    policy = DDPG.DDPG(state_dim, action_dim, max_action, device)  # , args.discount, args.tau)
    if args.generate_buffer:
        policy.load(f"./models/ddpg/behavioral_{setting}")

    # Initialize buffer
    replay_buffer = utils.ReplayBuffer(state_dim, action_dim, device)

    # Interact with the environment for max_timesteps
    for episode in range(args.episodes):
        state, _ = env.reset()
        episode_reward = 0
        step = 0
        while step < args.steps:
            step += 1
            # Select action with noise
            if (
                    (args.generate_buffer and np.random.uniform(0, 1) < args.rand_action_p) or
                    (args.train_behavioral and (args.steps * episode + step) < args.start_timesteps)
            ):
                # 产生数据集并且随机数小于设定的罪及操作的概率、或者训练behavioral时，
                action = env.action_space.sample()
            else:
                action = (
                        policy.select_action(np.array(state))
                        + np.random.normal(0, max_action * args.gaussian_std, size=action_dim)
                ).clip(-max_action, max_action)

            # Perform action
            next_state, reward, done, _, _ = env.step(action)
            reward = -reward ** 2 / 10
            done_bool = float(done)

            # Store data in replay buffer
            replay_buffer.add(state, action, next_state, reward, done_bool)

            state = next_state
            episode_reward += reward

            # Train agent after collecting sufficient data
            if args.train_behavioral:
                if (args.steps * episode + step) >= args.start_timesteps:
                    policy.train(replay_buffer, args.batch_size)
                summary_writer.add_scalar('episode_rewards', episode_reward, episode)
        if args.train_behavioral:
            print("Epoch/Episode: {}/{},reward: {}".format(episode + 1, args.episodes, episode_reward))
        else:
            print(f"Buffer add {args.steps * (episode + 1)} data")

    # Save final policy
    if args.train_behavioral:
        policy.save(f"./models/ddpg/behavioral_{setting}")

    # Save final buffer
    else:
        replay_buffer.save(f"./buffers/{buffer_name}")


# Trains BCQ offline
def train_BCQ(state_dim, action_dim, max_action, device, args):
    # For saving files
    setting = f"{args.env}_{args.seed}"
    buffer_name = f"{args.buffer_name}_{setting}"
    summary_writer = SummaryWriter(log_dir='./runs')

    # Initialize policy
    policy = BCQ.BCQ(state_dim, action_dim, max_action, device, args.discount, args.tau, args.lmbda, args.phi)

    # Load buffer
    replay_buffer = utils.ReplayBuffer(state_dim, action_dim, device)
    replay_buffer.load(f"./buffers/{buffer_name}")

    # while training_iters < args.max_timesteps:
    for episode in range(args.episodes):
        policy.train(replay_buffer, iterations=args.steps, batch_size=args.batch_size)
        score = eval_policy(policy, args.env, episodes=10, steps=100)
        # np.save(f"./results/BCQ_{setting}", evaluations)
        summary_writer.add_scalar('test score', score, episode)
        policy.save_model("./models/")
        print(f"Episode/Episodes: {episode + 1}/{args.episodes} score: {score}")


def eval_policy(policy, env_name, episodes=10, steps=100):
    eval_env = gym.make(env_name)
    total_reward = 0
    for _ in range(episodes):
        state, _ = eval_env.reset()
        step = 0
        while step < steps:
            step += 1
            action = policy.select_action(np.array(state))
            state, reward, done, _, _ = eval_env.step(action)
            total_reward += reward
    total_reward /= (episodes * steps)
    # print("---------------------------------------")
    # print(f"avg_reward: {total_reward:.3f}")
    # print("---------------------------------------")
    return total_reward


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--env", default="Pendulum-v1")  # OpenAI gym environment name   Pendulum-v1
    parser.add_argument("--seed", default=0, type=int)  # Sets Gym, PyTorch and Numpy seeds
    parser.add_argument("--buffer_name", default="Robust")  # Prepends name to filename
    parser.add_argument("--episodes", default=400, type=int)  # train epoch
    parser.add_argument("--steps", default=500, type=int)  # Max steps of each episode
    parser.add_argument("--start_timesteps", default=50000,
                        type=int)  # Time steps initial random policy is used before training behavioral
    parser.add_argument("--rand_action_p", default=0.3,
                        type=float)  # Probability of selecting random action during batch generation
    parser.add_argument("--gaussian_std", default=0.1,
                        type=float)  # Std of Gaussian exploration noise (Set to 0.1 if DDPG trains poorly)
    parser.add_argument("--batch_size", default=64, type=int)  # Mini batch size for networks
    parser.add_argument("--discount", default=0.99)  # Discount factor
    parser.add_argument("--tau", default=0.005)  # Target network update rate
    parser.add_argument("--lmbda", default=0.75)  # Weighting for clipped double Q-learning in BCQ
    parser.add_argument("--phi", default=0.05)  # Max perturbation hyper-parameter for BCQ
    parser.add_argument("--train_behavioral", action="store_true", default=False)  # If true, train behavioral (DDPG)
    parser.add_argument("--generate_buffer", action="store_true", default=False)  # If true, generate buffer
    args = parser.parse_args()

    print("---------------------------------------")
    if args.train_behavioral:
        print(f"Setting: Training behavioral, Env: {args.env}, Seed: {args.seed}")
    elif args.generate_buffer:
        print(f"Setting: Generating buffer, Env: {args.env}, Seed: {args.seed}")
    else:
        print(f"Setting: Training BCQ, Env: {args.env}, Seed: {args.seed}")
    print("---------------------------------------")

    if args.train_behavioral and args.generate_buffer:
        print("Train_behavioral and generate_buffer cannot both be true.")
        exit()

    if not os.path.exists("./results"):
        os.makedirs("./results")

    if not os.path.exists("./models"):
        os.makedirs("./models")

    if not os.path.exists("./buffers"):
        os.makedirs("./buffers")

    env = gym.make(args.env)
    env.action_space.seed(args.seed)
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    max_action = float(env.action_space.high[0])

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if args.train_behavioral or args.generate_buffer:
        interact_with_environment(env, state_dim, action_dim, max_action, device, args)
    else:
        train_BCQ(state_dim, action_dim, max_action, device, args)
